from typing import Dict, List, Optional
from ..format_utils import *
from ..instances.EntityFormatter import EntityFormatter


class SenteceFormatter:

    @classmethod
    def format_sentece_with_conditions(cls, data_item: Dict):

        sentence = data_item["sentence"]
        relations = data_item["relations"]
        sentence = add_sentence_prefix(sentence)
        conditions = add_conditions_prefix(EntityFormatter.from_relations(relations))

        return sentence + " " + conditions

    @classmethod
    def format_sentence_with_tag(cls, data_item: Dict):
        sentence = list(data_item["sentence"])
        relations = data_item["relations"]

        for r in relations:
            head = r["head"]
            tail = r["tail"]
            head_name = head["name"]
            tail_name = tail["name"]
            head_pos = head["pos"]
            tail_pos = tail["pos"]
            sentence[head_pos[0] : head_pos[1]] = f"<head>{head_name}</head>"
            sentence[tail_pos[0] : tail_pos[1]] = f"<tail>{tail_name}</tail>"
        return add_sentence_prefix("".join(sentence))

    @classmethod
    def format_sentence_for_re(cls, data_item: Dict):
        return add_sentence_prefix(data_item["sentence"])
